import os
import re
from pathlib import Path
from typing import List

from langchain_core.documents import Document
from langchain_community.document_loaders import PyPDFLoader, DirectoryLoader, TextLoader

from app.core.logger import logger
from app.rag.metadata_enhancer import enhance_document_metadata


def load_all_pdfs(pdf_dir: Path = None) -> List[Document]:
    """加载所有 PDF 文件并增强元数据"""
    if pdf_dir is None:
        from app.core.settings import settings
        pdf_dir = Path(settings.PDF_TXT_FILE_ROOT)
        
    logger.info(f"开始加载 PDF 文件，目录: {pdf_dir}")
    loader = DirectoryLoader(
        str(pdf_dir),
        glob="**/*.pdf",
        loader_cls=PyPDFLoader,
        show_progress=True,
        use_multithreading=True
    )
    docs = loader.load()
    
    # 增强每个文档的元数据
    enhanced_docs = []
    for i, doc in enumerate(docs):
        try:
            enhanced_doc = enhance_document_metadata(doc)
            enhanced_docs.append(enhanced_doc)
        except Exception as e:
            logger.warning(f"增强第 {i+1} 个 PDF 文档元数据失败: {e}")
            enhanced_docs.append(doc)  # 如果增强失败，保留原始文档
    
    logger.info(f"PDF 文件加载完成，共加载 {len(enhanced_docs)} 个文档")
    return enhanced_docs


def load_all_txts(txt_dir: Path = None) -> List[Document]:
    """
    加载所有TXT文件，从文件的前几行中提取元数据信息
    """
    if txt_dir is None:
        from app.core.settings import settings
        txt_dir = Path(settings.PDF_TXT_FILE_ROOT)
    
    if not txt_dir.exists():
        raise FileNotFoundError(f"TXT目录不存在: {txt_dir}")
    
    logger.info(f"开始加载TXT文件: {txt_dir}")
    documents = []
    success_count = 0
    error_count = 0
    
    for txt_file in txt_dir.glob("**/*.txt"):
        try:
            with open(txt_file, 'r', encoding='utf-8') as f:
                content = f.read().strip()
            if not content:
                logger.warning(f"跳过空文件: {txt_file}")
                continue

            # 读取前4行作为元数据
            with open(txt_file, 'r', encoding='utf-8') as f:
                first_four_lines = [f.readline().strip() for _ in range(4)]
            
            # 解析前4行信息
            metadata_info = {}
            for line in first_four_lines:
                if line.startswith("源文件路径:"):
                    metadata_info["source"] = line.split(":", 1)[1].strip()
                elif line.startswith("公司名:"):
                    metadata_info["company_name"] = line.split(":", 1)[1].strip()
                elif line.startswith("文件类型:"):
                    metadata_info["file_type"] = line.split(":", 1)[1].strip()
                elif line.startswith("总页数:"):
                    try:
                        metadata_info["total_pages"] = int(line.split(":", 1)[1].strip())
                    except ValueError:
                        pass  # 如果转换失败，跳过该字段

            # 创建 LangChain Document 对象
            doc = Document(
                page_content=content,
                metadata={
                    "source": str(txt_file.absolute()),
                    "file_name": txt_file.name,
                    "file_type": "txt",
                    **metadata_info  # 合并解析的元数据
                }
            )
            
            # 确保所有关键元数据字段都存在
            key_fields = ['source', 'file_name', 'company_name', 'file_type']
            for field in key_fields:
                if field not in doc.metadata or not doc.metadata[field]:
                    if field == 'source':
                        doc.metadata[field] = str(txt_file.absolute())
                    elif field == 'file_name':
                        doc.metadata[field] = txt_file.name
                    elif field == 'file_type':
                        doc.metadata[field] = doc.metadata.get('file_type', 'txt')
                    else:
                        doc.metadata[field] = 'unknown'
            
            # 增强元数据
            enhanced_doc = enhance_document_metadata(doc)
            documents.append(enhanced_doc)
            success_count += 1
            
            # 每处理10个文件记录一次进度
            if success_count % 10 == 0:
                logger.info(f"已处理 {success_count} 个 TXT 文件")
                
        except Exception as e:
            logger.error(f"读取文件失败: {txt_file}, 错误: {e}")
            error_count += 1

    logger.info(f"TXT 文件加载完成，成功: {success_count}, 失败: {error_count}, 总计: {len(documents)} 个文档")
    return documents


def load_all_documents():
    """
    一键加载 PDF + TXT 文件
    """
    docs = []

    # 加载 PDF
    pdf_docs = load_all_pdfs()
    docs.extend(pdf_docs)

    # 加载 TXT
    try:
        txt_docs = load_all_txts()
        docs.extend(txt_docs)
    except FileNotFoundError:
        logger.warning("TXT 目录未找到，跳过 TXT 文件加载")

    return docs